import pandas as pd
import numpy as np
from scipy import stats
import tushare as ts
from datetime import datetime,timedelta

class MultiFactorStrategy:
    def __init__(self, ts_token):
        """初始化多因子策略"""
        ts.set_token(ts_token)
        self.pro = ts.pro_api()
        
    def calculate_factors(self, stock_pool, 
                         value_weight=0.3,
                         momentum_weight=0.2,
                         quality_weight=0.3,
                         growth_weight=0.2):
        """计算多个因子并进行综合评分"""
        try:
            # 获取所有需要的数据
            factor_data = pd.DataFrame()
            
            # 获取基础数据
            factor_data = self._get_basic_data(stock_pool)
            if factor_data.empty:
                return None
                
            # 计算各类因子
            factor_data = self._calculate_value_factors(factor_data)    # 价值因子
            factor_data = self._calculate_momentum_factors(factor_data) # 动量因子
            factor_data = self._calculate_quality_factors(factor_data)  # 质量因子
            factor_data = self._calculate_growth_factors(factor_data)   # 成长因子
            
            # 标准化因子
            factor_cols = {
                'value_factors': ['pe_ttm', 'pb', 'ps_ttm'],
                'momentum_factors': ['momentum_20', 'momentum_60', 'momentum_120'],
                'quality_factors': ['roe', 'roa', 'gross_margin'],
                'growth_factors': ['revenue_growth', 'profit_growth']
            }
            
            # 对每个因子进行标准化
            for factor_group in factor_cols.values():
                for factor in factor_group:
                    if factor in factor_data.columns:
                        factor_data[f'{factor}_standardized'] = self._standardize(factor_data[factor])
            
            # 计算综合得分
            factor_data['final_score'] = (
                value_weight * self._calculate_group_score(factor_data, factor_cols['value_factors']) +
                momentum_weight * self._calculate_group_score(factor_data, factor_cols['momentum_factors']) +
                quality_weight * self._calculate_group_score(factor_data, factor_cols['quality_factors']) +
                growth_weight * self._calculate_group_score(factor_data, factor_cols['growth_factors'])
            )
            
            return factor_data
            
        except Exception as e:
            print(f"计算因子时出错: {str(e)}")
            return None

    def _get_basic_data(self, stock_pool):
        """获取基础数据"""
        try:
            # 获取最新交易日期
            latest_trade_date = self.pro.trade_cal(
                exchange='SSE',
                is_open=1,
                end_date=datetime.now().strftime('%Y%m%d'),
                fields='cal_date'
            ).iloc[-1]['cal_date']

            # 获取最新的每日指标
            daily_basic = self.pro.daily_basic(
                ts_code=','.join(stock_pool),
                trade_date=latest_trade_date,
                fields='ts_code,close,pe_ttm,pb,ps_ttm,total_mv,circ_mv,turnover_rate'
            )

            # 获取最新财务指标
            latest_quarter = self._get_latest_quarter()
            financial = self.pro.financial_indicator(
                ts_code=','.join(stock_pool),
                period=latest_quarter,
                fields='ts_code,roe,roa,grossprofit_margin,debt_to_assets,operating_revenue_yoy,basic_eps_yoy'
            )

            # 获取最近N天的日线数据用于计算动量因子
            end_date = latest_trade_date
            start_date = (datetime.strptime(latest_trade_date, '%Y%m%d') - timedelta(days=180)).strftime('%Y%m%d')
            
            daily_prices = self.pro.daily(
                ts_code=','.join(stock_pool),
                start_date=start_date,
                end_date=end_date,
                fields='ts_code,trade_date,close'
            )

            # 合并数据
            merged_data = pd.merge(daily_basic, financial, on='ts_code', how='left')
            
            # 添加价格数据到对象属性中，供动量因子计算使用
            self.price_data = daily_prices
            
            return merged_data
            
        except Exception as e:
            print(f"获取基础数据时出错: {str(e)}")
            return pd.DataFrame()

#     def _get_basic_data(self, stock_pool):
#         """获取基础数据"""
#         try:
#             # 获取最新的每日指标
#             daily_basic = self.pro.daily_basic(
#                 ts_code=','.join(stock_pool),
#                 trade_date=self.pro.trade_cal(is_open=1)['cal_date'].max()
#             )
            
#             # 获取最新的财务指标
#             financial = self.pro.financial_indicator(
#                 ts_code=','.join(stock_pool),
#                 period=self._get_latest_quarter()
#             )
            
#             # 合并数据
#             merged_data = pd.merge(daily_basic, financial, on='ts_code', how='left')
#             return merged_data
            
#         except Exception as e:
#             print(f"获取基础数据时出错: {str(e)}")
#             return pd.DataFrame()

    def _calculate_value_factors(self, df):
        """计算价值因子"""
        try:
            # 市盈率倒数 (收益率)
            df['ep_ratio'] = 1 / df['pe_ttm']
            
            # 市净率倒数 (账面价值率)
            df['bp_ratio'] = 1 / df['pb']
            
            # 市销率倒数 (销售收入率)
            df['sp_ratio'] = 1 / df['ps_ttm']
            
            return df
            
        except Exception as e:
            print(f"计算价值因子时出错: {str(e)}")
            return df


    def _calculate_momentum(self, stock_codes, period):
        """计算动量因子"""
        try:
            momentum = pd.DataFrame()
            
            # 使用之前存储的价格数据
            if not hasattr(self, 'price_data'):
                return pd.Series(index=stock_codes)
                
            price_data = self.price_data.copy()
            price_data['trade_date'] = pd.to_datetime(price_data['trade_date'])
            
            # 对每只股票计算动量
            for stock in stock_codes:
                stock_prices = price_data[price_data['ts_code'] == stock].sort_values('trade_date')
                if len(stock_prices) >= period:
                    # 计算期间收益率
                    latest_price = stock_prices['close'].iloc[-1]
                    past_price = stock_prices['close'].iloc[-period]
                    momentum.loc[stock, f'momentum_{period}'] = (latest_price / past_price) - 1
                    
            return momentum[f'momentum_{period}']
            
        except Exception as e:
            print(f"计算动量因子时出错: {str(e)}")
            return pd.Series()

    def _calculate_momentum_factors(self, df):
        """计算动量因子"""
        try:
            # 计算不同时间窗口的动量
            for period in [20, 60, 120]:
                df[f'momentum_{period}'] = self._calculate_momentum(df['ts_code'], period)
                
            return df
            
        except Exception as e:
            print(f"计算动量因子时出错: {str(e)}")
            return df

    def _calculate_quality_factors(self, df):
        """计算质量因子"""
        try:
            # ROE (净资产收益率)
            df['roe_factor'] = df['roe']
            
            # 毛利率
            df['gross_margin_factor'] = df['grossprofit_margin']
            
            # 资产周转率
            df['asset_turnover'] = df['total_revenue'] / df['total_assets']
            
            return df
            
        except Exception as e:
            print(f"计算质量因子时出错: {str(e)}")
            return df

    def _calculate_growth_factors(self, df):
        """计算成长因子"""
        try:
            # 营收增长率
            df['revenue_growth'] = df['revenue_growth']
            
            # 净利润增长率
            df['profit_growth'] = df['profit_growth']
            
            return df
            
        except Exception as e:
            print(f"计算成长因子时出错: {str(e)}")
            return df


    def _get_hs300_stocks(self):
        """获取沪深300成分股"""
        try:
            # 获取沪深300成分股
            df = self.pro.index_weight(
                index_code='399300.SZ',
                trade_date=self.pro.trade_cal(
                    exchange='SSE',
                    is_open=1,
                    end_date=datetime.now().strftime('%Y%m%d'),
                    fields='cal_date'
                ).iloc[-1]['cal_date']
            )
            return df['con_code'].tolist()
        except Exception as e:
            print(f"获取沪深300成分股时出错: {str(e)}")
            return []

    def select_stocks(self, stock_num=20)->pd.DataFrame:
        """选股主函数"""
        try:
            # 获取股票池（这里使用沪深300成分股作为示例）
            stock_pool = self._get_hs300_stocks()
            
            # 计算因子值
            factor_data = self.calculate_factors(stock_pool)
            if factor_data is None:
                return pd.DataFrame()
            
            # 根据综合得分排序并选出前N只股票
            selected_stocks = factor_data.nlargest(stock_num, 'final_score')
            
            return selected_stocks[['ts_code', 'final_score', 'pe_ttm', 'pb', 'roe']]
            
        except Exception as e:
            print(f"选股过程中出错: {str(e)}")
            return pd.DataFrame()

    def _standardize(self, series):
        """标准化处理"""
        return (series - series.mean()) / series.std()

    def _calculate_group_score(self, df, factors):
        """计算因子组合得分"""
        standardized_factors = [f'{factor}_standardized' for factor in factors]
        return df[standardized_factors].mean(axis=1)

    def _get_latest_quarter(self):
        """获取最新财报季度"""
        now = pd.Timestamp.now()
        year = now.year
        month = now.month
        if month < 4:
            return f"{year-1}0930"
        elif month < 7:
            return f"{year-1}1231"
        elif month < 10:
            return f"{year}0331"
        else:
            return f"{year}0630"
